from transformers import AutoModelForCausalLM, AutoTokenizer
import json
import pandas as pd

from vllm import LLM, SamplingParams
import re


zero_shot_sys_prompt="""role:Analyze provided power flow results to identify violated safety constraints.
Input: Bus voltage magnitude and angle

Each bus in the grid, node in the graph, contains 2 features:
- Bus: bus name, string
- Vm: Voltage Magnitude (p.u.)
- Va: Voltage Angle (degree)

The bus voltages will be given in columns "Bus | Vm | Va".
Voltage constraint will be given in interval such as [0.95 p.u., 1.05 p.u.]

Output json format for voltage constraint:
{"constraints":{"Voltage":{"buses":[buses that violate voltage constraint],"description":<Violation description>}}}

If none of constraint is violated, this should be output:
{"constraints":{"Voltage":{"buses":null,"description":<any reasonable description>}}}
"""
one_shot_sys_prompt=zero_shot_sys_prompt+"""Some examples are given below.
Example 1:
Bus Voltage Table(Bus | Vm | Va)
0 | 1 | -0.026
1 | 0.88 | -0.026
2 | 1 | -0.026
3 | 1.09 | -0.026
4 | 0.995 | -0.874
Voltage constraint:[0.93 p.u., 1.07 p.u.]

Output: {"constraints":{"Voltage":{"buses":[1,3],"description":"Bus 1 voltage (0.88 p.u.) is below minimum limit (0.93 p.u.). Bus 3 voltage (1.09 p.u.) exceeds maximum limit (1.07 p.u.)."}}}
"""
few_shot_sys_prompt=one_shot_sys_prompt+"""
Example 2:
Bus Voltage Table(Bus | Vm | Va)
0 | 1 | -0.026
1 | 0.94 | -0.026
2 | 1 | -0.026
3 | 1.07 | -0.026
4 | 0.995 | -0.874
Voltage constraint:[0.93 p.u., 1.07 p.u.]

Output:{"constraints":{"Voltage":{"buses":null,"description":"All bus voltages are within the acceptable range (0.93 p.u. to 1.07 p.u.)."}}}
"""


class VoltageConstraintChecker:

    def __init__(self, llm_model):
        self.llm = llm_model
        self.sampling_params = SamplingParams(temperature=0.8,top_k=3,max_tokens=8192)

    def read_voltage_data(self, file_path):
        """从Excel读取电压数据"""
        df = pd.read_excel(file_path)
        required_columns = ['Bus', 'Vm', 'Va']
        if not all(col in df.columns for col in required_columns):
            raise ValueError(f"Excel文件中必须包含以下列: {required_columns}")
        return df[required_columns]

    def format_prompt(self, sys_prompt,voltage_df):
        """生成两种电压约束的prompt"""
        voltage_table = "Bus Voltage Table(Bus | Vm | Va)\n"
        for _, row in voltage_df.iterrows():
            voltage_table += f"{row['Bus']} | {row['Vm']} | {row['Va']}\n"

        constraints = [[0.93, 1.07], [0.98, 1.02]]

        prompts = [
            {
                "messages": [
                    {"role": "system", "content": sys_prompt},
                    {
                        "role": "user",
                        "content": f"Analyze the voltage constraint violations of following result as required and output the results in JSON format.\n{voltage_table}\nVoltage constraint:[{c[0]} p.u., {c[1]} p.u.]",
                    },
                ]
            }
            for c in constraints
        ]

        return prompts

    def query_model(self, input_messages):
        """使用vLLM批量查询模型"""
        tokenizer = self.llm.get_tokenizer()
        prompts = [
            tokenizer.apply_chat_template(
                msg["messages"][:-1], tokenize=False, add_generation_prompt=True,enable_thinking=False
            )
            for msg in input_messages
        ]
        return self.llm.generate(prompts, self.sampling_params)

    def parse_response(self, response):
        """解析模型响应，提取违规信息
        使用正则表达式提取有问题，先直接解析JSON字符串"""
        # 提取JSON部分的正则表达式
        #pattern = r'{[\n\s]*\"constraints\":[\n\s]*({.*})}'
        #match = re.search(pattern, response, re.DOTALL)
        #if not match:
        #    return {"error": "未找到有效的约束信息"}

        try:
            # 简单验证是否为有效JSON
            #json_str =  match.group(0)
            #return json.loads(json_str)
            return json.loads(response)
        except json.JSONDecodeError as e:
            return {"error": f"JSON解析失败: {str(e)}"}
        except Exception as e:
            return {"error": f"响应格式解析失败: {str(e)}"}

    def run_test(self, excel_path):
        """完整测试流程"""
        # 1. 读取数据
        # voltage_data = self.read_voltage_data(excel_path)
        # 假设已经从Excel读取数据
        voltage_data = pd.DataFrame(
            {
                "Bus": [0, 1, 2, 3, 4],
                "Vm": [1.0, 0.94, 1.0, 1.07, 0.995],
                "Va": [-0.026, -0.026, -0.026, -0.026, -0.874],
            }
        )

        # 2. 生成prompt
        prompts = self.format_prompt(zero_shot_sys_prompt,voltage_data)

        # 3. 查询模型
        outputs = self.query_model(prompts)

        # 4. 解析结果
        results = []
        for i, output in enumerate(outputs):
            response = output.outputs[0].text
            parsed = self.parse_response(response)

            results.append({
                "constraint_range": [0.93, 1.07] if i == 0 else [0.98, 1.02],
                "raw_response": response,
                "parsed_result": parsed
            })

        return results

def run_zero_shot_test(model_dir,voltage_dfs:pd.DataFrame):
        
        llm_model=LLM(
            model=model_dir,
            tokenizer=model_dir,
            tensor_parallel_size=2,
            dtype="bfloat16",
            max_model_len=10240,
            gpu_memory_utilization=0.8,
        )
        checker = VoltageConstraintChecker(llm_model)
        results_mat = []
        for voltage_data in voltage_dfs:
            # 1. 生成prompt
            prompts = checker.format_prompt(zero_shot_sys_prompt,voltage_data)

            # 2. 查询模型
            outputs = checker.query_model(prompts)

            # 3. 解析结果
            results = []
            for i, output in enumerate(outputs):
                response = output.outputs[0].text
                parsed = checker.parse_response(response)

                results.append({
                    "constraint_range": [0.93, 1.07] if i == 0 else [0.98, 1.02],
                    "raw_response": response,
                    "parsed_result": parsed
                })
            results_mat.append(results)

        return results

# 使用示例
if __name__ == "__main__":
    model_dir = "/media/ubuntu/storage/qwen-3-30b"
    llm_model=LLM(
        model=model_dir,
        tokenizer=model_dir,
        tensor_parallel_size=2,
        dtype="bfloat16",
        max_model_len=10240,
        gpu_memory_utilization=0.8,
    )
    checker = VoltageConstraintChecker(llm_model)

    try:
        test_results = checker.run_test("voltage_data.xlsx")

        for i, result in enumerate(test_results, 1):
            print(f"测试 {i} (约束范围: {result['constraint_range']}):")
            print(f"模型原始响应:\n{result['raw_response']}")
            print(f"解析结果:\n{result['parsed_result']}\n")

    except Exception as e:
        print(f"测试过程中出错: {e}")
